import ipaddress
import socket
import requests
from typing import List, Set, Optional, Callable
import logging
import concurrent.futures

class NetworkScanner:
    """网络扫描器"""
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.custom_networks: Set[str] = set()
        self.port = 11434  # Ollama默认端口
    
    def add_network(self, network: str) -> bool:
        """添加自定义网段"""
        try:
            ipaddress.IPv4Network(network)
            self.custom_networks.add(network)
            return True
        except ValueError:
            self.logger.error(f"无效的网络格式: {network}")
            return False
    
    def get_local_networks(self) -> List[str]:
        """获取本地网段"""
        networks = []
        try:
            # 获取本机IP
            hostname = socket.gethostname()
            local_ip = socket.gethostbyname(hostname)
            
            # 添加本地网段
            network = ipaddress.IPv4Network(
                f"{local_ip}/24",
                strict=False
            )
            networks.append(str(network))
            
        except Exception as e:
            self.logger.error(f"获取本地网络失败: {str(e)}")
        
        return networks
    
    def scan_network(self, network: str, callback: Optional[Callable] = None) -> List[str]:
        """扫描网段"""
        servers = []
        try:
            network_obj = ipaddress.IPv4Network(network)
            
            with concurrent.futures.ThreadPoolExecutor(max_workers=50) as executor:
                future_to_ip = {
                    executor.submit(self.scan_port, str(ip), callback): str(ip)
                    for ip in network_obj.hosts()
                }
                
                for future in concurrent.futures.as_completed(future_to_ip):
                    ip = future_to_ip[future]
                    try:
                        result = future.result()
                        if result:
                            servers.append(result)
                    except Exception as e:
                        self.logger.debug(f"扫描 {ip} 失败: {str(e)}")
            
        except Exception as e:
            self.logger.error(f"扫描网段失败 {network}: {str(e)}")
        
        return servers
    
    def scan_port(self, ip: str, callback: Optional[Callable] = None) -> Optional[str]:
        """扫描单个IP的端口"""
        try:
            response = requests.get(
                f"http://{ip}:{self.port}/api/version",
                timeout=1
            )
            if response.status_code == 200:
                if callback:
                    callback(f"发现 Ollama 服务器: {ip}")
                return ip
        except:
            pass
        return None 